import heapq
from collections import defaultdict

def compressMethod(lsbs: list[str])->list[bytes,int,list]:
    """
    使用霍夫曼编码压缩二进制序列
    :param lsbs: 字符 '0' 和 '1' 组成的列表
    :return: 压缩后的字节数据, 原始长度, 码本
    """
    if not lsbs:
        raise ValueError("输入序列为空")

    lsbsStr = ''.join(lsbs)
    original_length = len(lsbsStr)  # 记录原始比特数

    # === 1. 统计频率 ===
    freq = defaultdict(int)
    for bit in lsbsStr:
        freq[bit] += 1

    # === 2. 构建霍夫曼树并生成码本 ===
    if len(freq) == 1:
        # 特殊情况：只有一种字符
        codebook = {list(freq.keys())[0]: '0'}
    else:
        heap = [(f, ch) for ch, f in freq.items()]
        heapq.heapify(heap)

        # 合并节点
        while len(heap) > 1:
            f1, c1 = heapq.heappop(heap)
            f2, c2 = heapq.heappop(heap)
            new_freq = f1 + f2
            new_code = None  # 内部节点无字符
            heapq.heappush(heap, (new_freq, (c1, c2, f1, f2)))  # 存储子树结构

        # 从最终树生成码本
        codebook = {}

        def build_code(node, prefix=""):
            if isinstance(node, str):  # 叶子节点
                codebook[node] = prefix
            else:
                left, right, f1, f2 = node
                build_code(left, prefix + "0")
                build_code(right, prefix + "1")

        final_node = heap[0][1]
        build_code(final_node)

    # === 3. 编码比特流 ===
    encoded_bits = ''.join(codebook[bit] for bit in lsbsStr)

    # === 4. 打包成字节流（补齐到8位）===
    padded_encoded = encoded_bits + '0' * ((8 - len(encoded_bits) % 8) % 8)
    byte_array = bytearray()
    for i in range(0, len(padded_encoded), 8):
        byte = padded_encoded[i:i+8]
        byte_array.append(int(byte, 2))
    compressed_bytes = bytes(byte_array)

    # === 5. 打包输出：压缩数据 + 码本 + 原始长度（用于解压）===
    # 这里我们简单用字符串表示码本（实际可用 pickle 或 JSON）
    import json
    codebook_data = json.dumps(codebook).encode('utf-8')
    # 原始长度用 4 字节无符号整数表示
    length_bytes = original_length.to_bytes(4, 'big')

    # 最终输出格式：[码本长度(2B)][码本][长度信息(4B)][压缩数据]
    cb_len = len(codebook_data)
    header = cb_len.to_bytes(2, 'big')  # 最大 65535
    final_compressed = header + codebook_data + length_bytes + compressed_bytes

    # 输出统计
    original_size = (original_length + 7) // 8
    compressed_size = len(final_compressed)
    ratio = compressed_size / original_size if original_size > 0 else 0

    # print(f"Huffman 压缩前: {original_size} 字节")
    # print(f"Huffman 压缩后: {compressed_size} 字节")
    # print(f"压缩比: {ratio:.2f}")
    # print(f"霍夫曼码本: {codebook}")

    return final_compressed, original_length, codebook


def decompressWithHuffman(compressed_data: bytes)-> list[str]:
    """
    解压霍夫曼压缩的数据
    :param compressed_data: 压缩后的字节数据（含头信息）
    :return: 解压后的二进制字符串列表
    """
    if len(compressed_data) < 6:
        raise ValueError("压缩数据过短")

    # === 1. 解析头部 ===
    cb_len = int.from_bytes(compressed_data[0:2], 'big')
    codebook_json = compressed_data[2:2+cb_len]
    original_length_bytes = compressed_data[2+cb_len:6+cb_len]
    original_length = int.from_bytes(original_length_bytes, 'big')
    huffman_bytes = compressed_data[6+cb_len:]

    # 恢复码本
    import json
    codebook = json.loads(codebook_json.decode('utf-8'))
    reverse_codebook = {v: k for k, v in codebook.items()}

    # === 2. 将字节转为比特流 ===
    bits = ''.join(f"{byte:08b}" for byte in huffman_bytes)

    # 截取实际编码部分（去除补位）
    total_encoded_bits = 0
    current = ""
    decoded = []

    for bit in bits:
        if len(decoded) >= original_length:
            break
        current += bit
        if current in reverse_codebook:
            decoded.append(reverse_codebook[current])
            current = ""
            total_encoded_bits += len(current)  # 实际解码长度

    # 截取正确长度
    result_str = ''.join(decoded)[:original_length]
    return list(result_str)  # 返回 list[str]


# === 示例使用 ===
if __name__ == "__main__":
    # 示例：0 占 80%，1 占 20%
    lsbs = ['0'] * 800 + ['1'] * 200

    print("=== 开始压缩 ===")
    compressed, orig_len, codebook = compressMethod(lsbs)

    print("\n=== 开始解压 ===")
    decompressed_lsbs = decompressWithHuffman(compressed)

    # 验证
    if decompressed_lsbs == lsbs:
        print("✅ 解压成功，数据一致！")
    else:
        print("❌ 解压失败，数据不一致！")
        print("原始:", lsbs[:50])
        print("解压:", decompressed_lsbs[:50])